package org.robotfish.learn;

import org.ejml.simple.SimpleMatrix;
import org.robotfish.geometry.TransformPoints;
import org.robotfish.optimize.MDFunction;
import org.robotfish.optimize.MDMinimizationDownhillSimplex;

public class EmpiricalHydroModel {
	double dt;
	double cycleLength;
	double regularization;
	int numKernelsTime;
	int numKernelsK;
	
	double kernelWidthTime, kernelWidthK;
	double kernelStartK, kernelStopK;
	
	int currentTimeStep;
	
	LinearModel FT, FN, T;
	
	int nData;
	

	public EmpiricalHydroModel(double dt, double cycleLength, double regularization, int numKernelsTime,  int numKernelsK,
			double kernelWidthTime, double kernelWidthK,
			double kernelStartK, double kernelStopK) {
		this.dt = dt;
		this.cycleLength  = cycleLength;
		this.regularization = regularization;
		this.numKernelsTime = numKernelsTime;
		this.numKernelsK = numKernelsK;
		this.currentTimeStep = -1;
		this.nData = 0;
		
		this.kernelWidthTime = kernelWidthTime;
		this.kernelWidthK = kernelWidthK;
		this.kernelStartK = kernelStartK;
		this.kernelStopK  = kernelStopK;
		
		setupModels();
	}
	
	/**
	 * Creates and initializes models for FT, FN, T. 
	 */
	public void setupModels() {
		// inputs are: v_T, v_N, w, k, t
		FT = new LinearModel(5, regularization);
		FN = new LinearModel(5, regularization);
		T  = new LinearModel(5, regularization);
		
		// FT = Gauss(t) + sign(v_T) v_T^2 + Gauss(k) w^2
		// Gauss(t)
		FT.addFeature( new GaussianFeatures(
				sequence(0, dt, numKernelsTime, false),  // kernels
				kernelWidthTime,    // width
				4,                  // time
				dt)                 // cycle length
		);
		// sign(v_T) v_T^2:
		FT.addFeature( new PolyFeatures(new int[]{0}, new int[]{2}, true, true) );
		// Gauss(k) w^2
		FT.addFeature( new GaussianProdFeatures(
				sequence(kernelStartK, kernelStopK, numKernelsK, true), 
				kernelWidthK, 
				3,  // index of k (Gaussian feature)
				0,  // no cycle
				2,  // index of w (polynomial feature)
				2   // power of w
		));

		/////////////////////////////////////////////
		// FN = Gauss(t) + sign(v_N) v_N^2
		// Gauss(t)
		FN.addFeature( new GaussianFeatures(
				sequence(0, dt, numKernelsTime, false),  // kernels
				kernelWidthTime,    // width
				4,                  // time
				dt)                 // cycle length
		);
		// sign(v_N) v_N^2:
		FN.addFeature( new PolyFeatures(new int[]{1}, new int[]{2}, true, true) );
		
		
		//////////////////////////////////////////////
		// T = Gauss(t) + sign(w) w^2 + Gauss(k) v_T^2
		// Gauss(t)
		T.addFeature( new GaussianFeatures(
				sequence(0, dt, numKernelsTime, false),  // kernels
				kernelWidthTime,    // width
				4,                  // time
				dt)                 // cycle length
		);
		// sign(w) w^2:
		T.addFeature( new PolyFeatures(new int[]{2}, new int[]{2}, true, true) );
		// Gauss(k) w^2:
		T.addFeature( new GaussianProdFeatures(
				sequence(kernelStartK, kernelStopK, numKernelsK, true), 
				kernelWidthK, 
				3,  // index of k (Gaussian feature)
				0,  // no cycle
				0,  // index of v_T (polynomial feature)
				2   // power of w
		));
	}
	
	/**
	 * @param force     3-vector (x, y, angle)
	 * @param velocity  3-vector (x, y, angle)
	 * @param COG       3-vector (x, y, angle)
	 * @param k         curvature
	 * @param timeStep  time step
	 */
	public void learnHydroModel( SimpleMatrix force, SimpleMatrix velocity, SimpleMatrix COG, double k, int timeStep ) {
		if (currentTimeStep>=0) {
			// checking if time has been increased by just one, otherwise throwing exception:
			if (timeStep!=currentTimeStep+1) {
				throw new RuntimeException(String.format(
						"Timestep was not increased by one. Current timestep was %d and new was %d.", 
						currentTimeStep, timeStep) );
			}
		}
		currentTimeStep = timeStep;
		// column vector holding the input: (v_T, v_N, w, k, t)
		SimpleMatrix input = makeInput(velocity, COG, k, timeStep);
				
		SimpleMatrix forceRel  = TransformPoints.toRelative(force, COG.get(2));
		nData++;
		
		FT.addDataPoint( input, forceRel.get(0) );
		FN.addDataPoint( input, forceRel.get(1) );
		T.addDataPoint(  input, forceRel.get(2) );
		
		if (nData<=200) {
			FT.doRegFullRegression();
			FN.doRegFullRegression();
			T.doRegFullRegression();
		} else {
			// just updating model using the last point:
			FT.updateRegression();
			FN.updateRegression();
			T.updateRegression();
		}
	}

	/**
	 * 
	 * @param velocity 3-vector (v_x, v_y, w)
	 * @param COG      3-vector (x, y, alpha)
	 * @param k        curvature
	 * @param timeStep
	 * @return input of (v_T, v_N, w, k, t)
	 */
	public SimpleMatrix makeInput(SimpleMatrix velocity, SimpleMatrix COG,
			double k, int timeStep) {
		SimpleMatrix input = new SimpleMatrix(5, 1);

		// rotating velocity and force by COG.get(2) clockwise:
		SimpleMatrix velocityRel = TransformPoints.toRelative(velocity, COG.get(2));
		
		input.insertIntoThis(0, 0, velocityRel);
		input.set( 3, k );                       // k
		input.set( 4, timeStep*dt );             // t = dt * timeStep
		return input;
	}

	/**
	 * @param velocity 3-vector of (vx, vy, w)
	 * @param COG      3-vector of (x, y, angle)
	 * @param k
	 * @param timeStep
	 * @return 3-vector of (FT, FN, T), in absolute coordinates.
	 */
	public SimpleMatrix predict(SimpleMatrix velocity, SimpleMatrix COG, double k, int timeStep) {
		SimpleMatrix input = makeInput(velocity, COG, k, timeStep);
		SimpleMatrix forceRel = new SimpleMatrix( new double[][] {
			{ FT.getPrediction(input) },
			{ FN.getPrediction(input) },
			{ T.getPrediction(input) }
		});
		// rotating back:
		return TransformPoints.toRelative( forceRel, -COG.get(2) );
	}
	
	/**
	 * 
	 * @param mechanicalModel
	 * @return 3-vector of speeds (vx, vy, w)
	 */
	public SimpleMatrix findVelocity( RigidLinkFishState state, MechanicalModel mechanicalModel, FishMotion fishMotion,
			SimpleMatrix initialGuess) {
		// using initialGuess:
		//SimpleMatrix velocityCOG = initialGuess;
		
		VelocityMinimizer minfun = new VelocityMinimizer(mechanicalModel, state, fishMotion);
		MDMinimizationDownhillSimplex minimizer =  new MDMinimizationDownhillSimplex (minfun);
		double init[] = {initialGuess.get(0), initialGuess.get(1), initialGuess.get(2)};
		minimizer.setSimplex(init, 0.4);
		minimizer.minimize();
		
		double[] minX = minimizer.getMinX();
		
		return new SimpleMatrix( 3, 1, true, minX );
	}
	
	/**
	 * Nested class for minimization.
	 */
	class VelocityMinimizer implements MDFunction {
		RigidLinkFishState state;
		MechanicalModel mechanicalModel;
		FishMotion fishMotion;
		
		public VelocityMinimizer(MechanicalModel m, RigidLinkFishState state, 
				FishMotion fishMotion) {
			this.state = state;
			this.mechanicalModel = m;
			this.fishMotion = fishMotion;
		}
		
		@Override
		public int argumentLength() {
			return 3;
		}
		@Override
		public double f(double[] x) {
			// velocityCOG is used for searching:
			SimpleMatrix velocityCOG = new SimpleMatrix(3, 1, true, x);
			RigidLinkFishState nextState = mechanicalModel.getNextState(state, velocityCOG, fishMotion);
			SimpleMatrix force = mechanicalModel.getForceCOG(nextState);
			SimpleMatrix cog   = mechanicalModel.getPositionCOG(nextState);
			
			double k = mechanicalModel.getCurvature(nextState);
			int timeStep = mechanicalModel.getTimeStep(nextState);
			// x is 3-vector of velocity:
			SimpleMatrix forceE = EmpiricalHydroModel.this.predict(velocityCOG, cog, k, timeStep);
			
			// find the difference and add up absolute values of all elements:
			SimpleMatrix diff = force.minus(forceE);
			double diffSum = 0;
			for (int i=0; i<3; i++) {
				diffSum += Math.abs( diff.get(i) );
			}
			return diffSum;
		}
	}
	
	/**
	 * 
	 * @param start  start of the sequence
	 * @param stop   end of the sequence
	 * @param num    number of elements in the sequence
	 * @param last   whether to include or exclude last point
	 * @return sequence with fixed increment:
	 * <li>If last is included: {@code (stop-start) / (num-1) }
	 * <li>If last in not included: {@code (stop-start) / num }
	 */
	public static double[] sequence(double start, double stop, int num, boolean last) {
		double[] seq = new double[num];
		double diff;
		if (last) {
			diff = (stop-start) / (num-1);
		} else {
			diff = (stop-start) / num;
		}
		for (int i=0; i<num; i++) {
			seq[i] = start + i*diff;
		}
		return seq;
	}

	public double getCycleLength() {
		return cycleLength;
	}

	public void setCycleLength(double cycleLength) {
		this.cycleLength = cycleLength;
	}

	public double getRegularization() {
		return regularization;
	}

	public void setRegularization(double regularization) {
		this.regularization = regularization;
	}

	public int getNumKernelsTime() {
		return numKernelsTime;
	}

	public void setNumKernelsTime(int numKernelsTime) {
		this.numKernelsTime = numKernelsTime;
	}

	public int getNumKernelsK() {
		return numKernelsK;
	}

	public void setNumKernelsK(int numKernelsK) {
		this.numKernelsK = numKernelsK;
	}

	public double getKernelWidthTime() {
		return kernelWidthTime;
	}

	public void setKernelWidthTime(double kernelWidthTime) {
		this.kernelWidthTime = kernelWidthTime;
	}

	public double getKernelWidthK() {
		return kernelWidthK;
	}

	public void setKernelWidthK(double kernelWidthK) {
		this.kernelWidthK = kernelWidthK;
	}

	public double getKernelStartK() {
		return kernelStartK;
	}

	public void setKernelStartK(double kernelStartK) {
		this.kernelStartK = kernelStartK;
	}

	public double getKernelStopK() {
		return kernelStopK;
	}

	public void setKernelStopK(double kernelStopK) {
		this.kernelStopK = kernelStopK;
	}

	public void setDt(double dt) {
		this.dt = dt;
	}

	public double getDt() {
		return dt;
	}

	public int getnData() {
		return nData;
	}

}
